
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "mbedtls/md.h"
#include "abcrypto.h"
#include "hash.h"

static int init_ctx(void **ctx, uint8_t *key, uint32_t klen, uint16_t alg)
{
	mbedtls_md_context_t *mbedctx = NULL;
	const mbedtls_md_info_t *md_info = NULL;

	mbedctx = malloc(sizeof(mbedtls_md_context_t));
	mbedtls_md_init(mbedctx);

	switch (alg & 0x0f) {
	case ABCRYPTO_AUTH_ALG_MD5:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_MD5);

		break;
	case ABCRYPTO_AUTH_ALG_SHA1:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA1);
		break;
	case ABCRYPTO_AUTH_ALG_SHA224:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA224);
		break;
	case ABCRYPTO_AUTH_ALG_SHA256:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA256);
		break;
	case ABCRYPTO_AUTH_ALG_SHA384:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA384);
		break;
	case ABCRYPTO_AUTH_ALG_SHA512:
		// 获取 SHA512 算法信息
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA512);
		break;
	case ABCRYPTO_AUTH_ALG_SHA3_224:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA3_224);
		break;
	case ABCRYPTO_AUTH_ALG_SHA3_256:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA3_256);
		break;
	case ABCRYPTO_AUTH_ALG_SHA3_384:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA3_384);
		break;
	case ABCRYPTO_AUTH_ALG_SHA3_512:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA3_512);
		break;
	default:
		break;
	}

	if (md_info == NULL) {
		fprintf(stderr,
			"Failed to get SHA512 algorithm information.\n");
		free(mbedctx);
		return -1;
	}

	// 设置消息摘要上下文
	int ret = mbedtls_md_setup(mbedctx, md_info, 0);
	if (ret != 0) {
		fprintf(stderr, "Failed to setup SHA512 context.\n");
		mbedtls_md_free(mbedctx);
		free(mbedctx);
		return -1;
	}

	// 开始哈希计算
	ret = mbedtls_md_starts(mbedctx);
	if (ret != 0) {
		fprintf(stderr, "Failed to start SHA512 calculation.\n");
		mbedtls_md_free(mbedctx);
		free(mbedctx);
		return 1;
	}

	*ctx = mbedctx;

	return 0;
}

static void free_ctx(void *ctx)
{
	mbedtls_md_free(ctx);
	free(ctx);
	return;
}

static int update(void *ctx, uint8_t *data, uint32_t len)
{
	return mbedtls_md_update(ctx, data, len);
}

static int finish(void *ctx, uint8_t *dgst)
{
	return mbedtls_md_finish(ctx, dgst);
}

static int digest(uint8_t *data, uint32_t len, uint8_t *key, uint32_t klen,
		  uint8_t *dgst, uint16_t alg)
{
	const mbedtls_md_info_t *md_info = NULL;

	switch (alg & 0x0f) {
	case ABCRYPTO_AUTH_ALG_MD5:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_MD5);

		break;
	case ABCRYPTO_AUTH_ALG_SHA1:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA1);
		break;
	case ABCRYPTO_AUTH_ALG_SHA224:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA224);
		break;
	case ABCRYPTO_AUTH_ALG_SHA256:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA256);
		break;
	case ABCRYPTO_AUTH_ALG_SHA384:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA384);
		break;
	case ABCRYPTO_AUTH_ALG_SHA512:
		// 获取 SHA512 算法信息
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA512);
		break;
	case ABCRYPTO_AUTH_ALG_SHA3_224:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA3_224);
		break;
	case ABCRYPTO_AUTH_ALG_SHA3_256:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA3_256);
		break;
	case ABCRYPTO_AUTH_ALG_SHA3_384:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA3_384);
		break;
	case ABCRYPTO_AUTH_ALG_SHA3_512:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA3_512);
		break;
	default:
		break;
	}

	if (md_info == NULL) {
		fprintf(stderr, "Failed to get algorithm info\n");
		return -1;
	}

	return mbedtls_md(md_info, data, len, dgst);
}

static struct abcrypto_auth hash = {
	.init_ctx = init_ctx,
	.free_ctx = free_ctx,
	.update = update,
	.finish = finish,
	.digest = digest,
};

struct abcrypto_auth *mbed_hash(void)
{
	return &hash;
}

static int hmac_init_ctx(void **ctx, uint8_t *key, uint32_t klen, uint16_t alg)
{
	mbedtls_md_context_t *mbedctx = NULL;
	const mbedtls_md_info_t *md_info = NULL;

	mbedctx = malloc(sizeof(mbedtls_md_context_t));
	if (NULL == mbedctx)
		return -1;

	mbedtls_md_init(mbedctx);

	switch (alg & 0x0f) {
	case ABCRYPTO_AUTH_ALG_MD5:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_MD5);
		break;
	case ABCRYPTO_AUTH_ALG_SHA1:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA1);
		break;
	case ABCRYPTO_AUTH_ALG_SHA224:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA224);
		break;
	case ABCRYPTO_AUTH_ALG_SHA256:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA256);
		break;
	case ABCRYPTO_AUTH_ALG_SHA384:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA384);
		break;
	case ABCRYPTO_AUTH_ALG_SHA512:
		// 获取 SHA512 算法信息
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA512);
		break;
	case ABCRYPTO_AUTH_ALG_SHA3_224:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA3_224);
		break;
	case ABCRYPTO_AUTH_ALG_SHA3_256:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA3_256);
		break;
	case ABCRYPTO_AUTH_ALG_SHA3_384:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA3_384);
		break;
	case ABCRYPTO_AUTH_ALG_SHA3_512:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA3_512);
		break;
	default:
		break;
	}

	if (md_info == NULL) {
		mbedtls_md_free(mbedctx);
		free(mbedctx);
		fprintf(stderr, "Failed to get SHA512 algorithm info\n");
		return -1;
	}

	// 初始化 HMAC 上下文
	if (mbedtls_md_setup(mbedctx, md_info, 1) != 0) {
		fprintf(stderr, "Failed to setup HMAC context\n");
		mbedtls_md_free(mbedctx);
		free(mbedctx);
		return -1;
	}

	// 初始化 HMAC 计算
	if (mbedtls_md_hmac_starts(mbedctx, key, klen) != 0) {
		fprintf(stderr, "Failed to start HMAC calculation\n");
		mbedtls_md_free(mbedctx);
		free(mbedctx);
		return -1;
	}

	*ctx = mbedctx;

	return 0;
}

static void hmac_free_ctx(void *ctx)
{
	mbedtls_md_free(ctx);
	free(ctx);
	return;
}

static int hmac_update(void *ctx, uint8_t *data, uint32_t len)
{
	return mbedtls_md_hmac_update(ctx, data, len);
}

static int hmac_finish(void *ctx, uint8_t *dgst)
{
	return mbedtls_md_hmac_finish(ctx, dgst);
}

static int hmac_digest(uint8_t *data, uint32_t len, uint8_t *key, uint32_t klen,
		       uint8_t *dgst, uint16_t alg)
{
	const mbedtls_md_info_t *md_info = NULL;

	switch (alg & 0x0f) {
	case ABCRYPTO_AUTH_ALG_MD5:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_MD5);

		break;
	case ABCRYPTO_AUTH_ALG_SHA1:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA1);
		break;
	case ABCRYPTO_AUTH_ALG_SHA224:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA224);
		break;
	case ABCRYPTO_AUTH_ALG_SHA256:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA256);
		break;
	case ABCRYPTO_AUTH_ALG_SHA384:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA384);
		break;
	case ABCRYPTO_AUTH_ALG_SHA512:
		// 获取 SHA512 算法信息
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA512);
		break;
	case ABCRYPTO_AUTH_ALG_SHA3_224:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA3_224);
		break;
	case ABCRYPTO_AUTH_ALG_SHA3_256:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA3_256);
		break;
	case ABCRYPTO_AUTH_ALG_SHA3_384:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA3_384);
		break;
	case ABCRYPTO_AUTH_ALG_SHA3_512:
		md_info = mbedtls_md_info_from_type(MBEDTLS_MD_SHA3_512);
		break;
	default:
		break;
	}

	if (md_info == NULL) {
		fprintf(stderr, "Failed to get algorithm info\n");
		return -1;
	}

	return mbedtls_md_hmac(md_info, key, klen, data, len, dgst);
}

static struct abcrypto_auth hmac = {
	.init_ctx = hmac_init_ctx,
	.free_ctx = hmac_free_ctx,
	.update = hmac_update,
	.finish = hmac_finish,
	.digest = hmac_digest,
};

struct abcrypto_auth *mbed_hmac(void)
{
	// return NULL;
	return &hmac;
}